# Perfect matching in complete bipartite graphs. 

abstract type PerfectBipartiteMatchingSolver end
function build!(solver::PerfectBipartiteMatchingSolver, reward::Matrix{Distribution})
  nothing
end

struct PerfectBipartiteMatching <: CombinatorialInstance
  # Mandatory properties. 
  n_arms::Int

  # Probability distributions for the arm rewards. 
  reward::Matrix{Distribution}
  optimal_reward::Float64 # Cache. 

  # Internal solver. 
  solver::PerfectBipartiteMatchingSolver

  function PerfectBipartiteMatching(reward::Matrix{Distribution}, solver::PerfectBipartiteMatchingSolver)
    n = size(reward, 1)
    if n != size(reward, 2)
      error("Graph is not bipartite complete: reward matrix must be square.")
    end

    # Prepare the solver to be used (if required). 
    build!(solver, reward)

    # Precompute certain results. 
    weights = Dict((i, j) => mean(reward[i, j]) for i in 1:n, j in 1:n)
    assignment = solve_linear(solver, weights)
    optimal_reward = sum(weights[(i, j)] for (i, j) in assignment)

    # Done! 
    return new(n ^ 2, reward, optimal_reward, solver)
  end
end

function initial_state(instance::PerfectBipartiteMatching) 
  n = size(instance.reward, 1)
  zeroCounts = Dict((i, j) => 0 for i in 1:n, j in 1:n)
  zeroRewards = Dict((i, j) => 0.0 for i in 1:n, j in 1:n)
  return State{Tuple{Int, Int}}(0, 0.0, 0.0, zeroCounts, zeroRewards, copy(zeroRewards))
end

function initial_trace(instance::PerfectBipartiteMatching) 
  return Trace{Tuple{Int, Int}}(State{Tuple{Int, Int}}[], Vector{Tuple{Int, Int}}[], Float64[], PolicyDetails[])
end

function pull(instance::PerfectBipartiteMatching, arms::Vector{Tuple{Int, Int}}) 
  return Float64[rand(instance.reward[arm[1], arm[2]]) for arm in arms]
end

solve_linear(instance::PerfectBipartiteMatching, weights::Dict{Tuple{Int, Int}, Float64}) = solve_linear(instance.solver, weights)

function is_feasible(instance::PerfectBipartiteMatching, arms::Vector{Tuple{Int, Int}}) 
  if length(arms) > instance.n_arms
    return false
  end

  # Check whether each node is taken at most once on each side. 
  lefts = sort([arm[1] for arm in arms])
  rights = sort([arm[2] for arm in arms])

  return lefts == unique(lefts) && rights == unique(rights)
end
